import math

import common.img_utils.fh.OptimizedUnionFind as uf


def get_diff(img, x1, y1, x2, y2):
    """Get the L2 distance between [y1, x1] and [y2, x2] in each color channel.

    Parameters
    ----------
    img : tuple
        tuple of the input image with each splitted channel (image band 1 (h x w), image band 2 (h x w), ...)
    x1 : int
        x coordinate of the first point
    y1 : int
        y coordinate of the first point
    x2 : int
        x coordinate of the second point
    y2 : int
        y coordinate of the second point

    Returns
    -------
    ndarray
        The summed L2 distance of the image
    """

    l2_dist_sum = 0
    range_num_for_loop = len(img)

    for i in range(range_num_for_loop):
        l2_dist_sum += (img[i][y1, x1] - img[i][y2, x2]) ** 2
    return math.sqrt(l2_dist_sum)


def get_threshold(k, size):
    return k / size


def vertex_id(x, y, width):
    return y * width + x


def create_edge(img, width, x1, y1, x2, y2):
    w = get_diff(img, x1, y1, x2, y2)
    return (vertex_id(x1, y1, width), vertex_id(x2, y2, width), w)


def build_graph(img, width, height):
    """Build a graph of input image tuple

    Parameters
    ----------
    img : tuple
        tuple of the input image with each splitted channel (image band 1 (h x w), image band 2 (h x w), ...)
    width : int
        the width of multi-channgel input image
    height : int
        the height of multi-channgel input image

    Returns
    -------
    graph : list
        a graph of input image tuple
    """
    graph = []

    for y in range(height):
        for x in range(width):
            if x < width - 1:
                graph.append(create_edge(img, width, x, y, x + 1, y))
            if y < height - 1:
                graph.append(create_edge(img, width, x, y, x, y + 1))
            if x < width - 1 and y < height - 1:
                graph.append(create_edge(img, width, x, y, x + 1, y + 1))
            if x < width - 1 and y > 0:
                graph.append(create_edge(img, width, x, y, x + 1, y - 1))

    return graph


def remove_small_component(ufset, sorted_graph, min_size):
    for edge in sorted_graph:
        u = ufset.find(edge[0])
        v = ufset.find(edge[1])

        if u != v:
            if ufset.size_of(u) < min_size or ufset.size_of(v) < min_size:
                ufset.merge(u, v)

    return ufset


def segment_graph(sorted_graph, num_node, k):
    ufset = uf.OptimizedUnionFind(num_node)
    threshold = [get_threshold(k, 1)] * num_node

    for edge in sorted_graph:
        u = ufset.find(edge[0])
        v = ufset.find(edge[1])
        w = edge[2]

        if u != v:
            if w <= threshold[u] and w <= threshold[v]:
                ufset.merge(u, v)
                parent = ufset.find(u)
                threshold[parent] = w + get_threshold(k, ufset.size_of(parent))

    return ufset
